#ifndef MISH_GRAD_N_D_H
#define MISH_GRAD_N_D_H

#include "kernel_operator.h"

namespace MishGrad {
using namespace AscendC;

constexpr int32_t DOUBLE_BUFFER_NUM = 1;
constexpr int32_t SINGLE_BUFFER_NUM = 1;

constexpr int32_t BYTE_BLOCK = 32;
constexpr int32_t TMP_BUF_COUNT = 2;

constexpr float POS_ONE = 1.0;
constexpr float FOUR = 4.0;
constexpr float NEG_TWO = -2.0;

template <typename T>
class MishGradND {
public:
    __aicore__ inline MishGradND() = default;
    __aicore__ inline void Init(GM_ADDR grad, GM_ADDR x, GM_ADDR y, const MishGradTilingData* __restrict tilingData);
    __aicore__ inline void Process();

private:
    template <typename T1, typename T2>
    __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
    };

    template <typename T1, typename T2>
    __aicore__ inline T1 CeilAlignA2B(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
    };

    __aicore__ inline void CopyIn(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void Compute(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void CopyOut(int64_t gmOffset, int64_t dataCount);

private:
    TPipe pipe;

    TQue<QuePosition::VECIN, SINGLE_BUFFER_NUM> gradQue;
    TQue<QuePosition::VECIN, DOUBLE_BUFFER_NUM> xQue;
    TQue<QuePosition::VECOUT, DOUBLE_BUFFER_NUM> yQue;

    TBuf<QuePosition::VECCALC> tempValBuf;
    GlobalTensor<T> gradGM, xGM, yGM;
    LocalTensor<T> t1, t2;
    int64_t blockIdx = 0;
    uint64_t perBlockCount = 0;

    // tiling params
    uint64_t ubMaxProcCount = 0;
    uint64_t totalDataCount = 0;
    uint64_t loopCount = 0;
    uint64_t tailCount = 0;
};

template <typename T>
__aicore__ inline void MishGradND<T>::Init(GM_ADDR grad, GM_ADDR x, GM_ADDR y,
                                           const MishGradTilingData* __restrict tilingData) {
    blockIdx = GetBlockIdx();
    perBlockCount = BYTE_BLOCK / sizeof(T);
    ubMaxProcCount = tilingData->ubMaxProcCount;
    totalDataCount = tilingData->totalDataCount;
    loopCount = tilingData->loopCount;
    tailCount = tilingData->tailCount;

    gradGM.SetGlobalBuffer((__gm__ T*)grad, totalDataCount);
    xGM.SetGlobalBuffer((__gm__ T*)x, totalDataCount);
    yGM.SetGlobalBuffer((__gm__ T*)y, totalDataCount);

    uint64_t singleBufferSize = ubMaxProcCount * sizeof(T);
    pipe.InitBuffer(gradQue, SINGLE_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(xQue, DOUBLE_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(yQue, DOUBLE_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(tempValBuf, singleBufferSize * TMP_BUF_COUNT);

    t1 = tempValBuf.Get<T>();
    t2 = t1[ubMaxProcCount];
}

template <typename T>
__aicore__ inline void MishGradND<T>::Process() {
    int64_t gmOffset = 0;
    for (int64_t i = 0; i < loopCount; i++) {
        CopyIn(gmOffset, ubMaxProcCount);
        Compute(gmOffset, ubMaxProcCount);
        CopyOut(gmOffset, ubMaxProcCount);
        gmOffset += ubMaxProcCount;
    }
    if (tailCount) {
        int64_t alignDataCount = CeilAlignA2B(tailCount, perBlockCount);
        CopyIn(gmOffset, alignDataCount);
        Compute(gmOffset, alignDataCount);
        CopyOut(gmOffset, alignDataCount);
    }
}

template <typename T>
__aicore__ inline void MishGradND<T>::CopyIn(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> gradInLT = gradQue.AllocTensor<T>();
    LocalTensor<T> xInLT = xQue.AllocTensor<T>();
    DataCopy(gradInLT, gradGM[gmOffset], dataCount);
    DataCopy(xInLT, xGM[gmOffset], dataCount);
    gradQue.EnQue(gradInLT);
    xQue.EnQue(xInLT);
}

template <typename T>
__aicore__ inline void MishGradND<T>::Compute(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> xLT = xQue.DeQue<T>();
    LocalTensor<T> yLT = yQue.AllocTensor<T>();

    Exp(t1, xLT, dataCount);              // exp_x = np.exp(input_x)
    Adds(t2, t1, T(POS_ONE), dataCount);  // exp_add_x = exp_x + 1
    // temp1 = exp_add_x * exp_x * input_x
    Mul(yLT, t2, t1, dataCount);
    Mul(xLT, yLT, xLT, dataCount);

    // rec_exp_add_x = 1.0 / (exp_add_x * exp_add_x + 1.0)
    Duplicate(yLT, T(POS_ONE), dataCount);
    Mul(t2, t2, t2, dataCount);
    Adds(t2, t2, T(POS_ONE), dataCount);
    Div(t2, yLT, t2, dataCount);

    // temp2 = temp1 + rec_exp_add_x * rec_exp_add_x * 4
    // result = temp2 - 2 * rec_exp_add_x + 1
    Mul(yLT, xLT, t2, dataCount);
    Mul(yLT, yLT, t2, dataCount);
    Muls(yLT, yLT, T(FOUR), dataCount);

    Muls(t2, t2, T(NEG_TWO), dataCount);
    Add(yLT, yLT, t2, dataCount);
    Adds(yLT, yLT, T(POS_ONE), dataCount);

    LocalTensor<T> gradLT = gradQue.DeQue<T>();
    Mul(yLT, yLT, gradLT, dataCount);

    yQue.EnQue(yLT);
    gradQue.FreeTensor(gradLT);
    xQue.FreeTensor(xLT);
}

template <typename T>
__aicore__ inline void MishGradND<T>::CopyOut(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> yOutLT = yQue.DeQue<T>();
    DataCopy(yGM[gmOffset], yOutLT, dataCount);
    yQue.FreeTensor(yOutLT);
}
}  // namespace MishGrad

#endif  // MISH_GRAD_N_D_H